package com.yeskery.nut.extend.mcp.server;

import com.yeskery.nut.bean.ApplicationContext;
import com.yeskery.nut.bean.FactoryBeanRegister;
import com.yeskery.nut.bind.FitValueHelper;
import com.yeskery.nut.extend.mcp.annotation.McpTool;
import com.yeskery.nut.extend.mcp.annotation.McpToolProperty;
import com.yeskery.nut.extend.mcp.core.Tool;
import com.yeskery.nut.extend.mcp.core.schema.ToolSchema;
import com.yeskery.nut.util.StringUtils;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Mcp工具bean后置处理器
 * @author sprout
 * 2025-04-18 11:22
 */
public class McpToolBeanPostProcessor extends BaseMcpBeanPostProcessor {

    /**
     * 构建Mcp工具bean后置处理器
     * @param mcpServerCapability mcp服务能力
     */
    public McpToolBeanPostProcessor(McpServerCapability mcpServerCapability) {
        super(mcpServerCapability);
    }

    @Override
    protected Class<? extends Annotation> getMethodTargetAnnotation() {
        return McpTool.class;
    }

    @Override
    protected void doProcess(Object bean, Class<?> beanClass, Method method, ApplicationContext applicationContext, FactoryBeanRegister factoryBeanRegister) {
        Tool tool = createTool(method);
        ToolSchema.Tool schemaTool = tool.getTool();
        Map<String, ToolSchema.Property> properties = schemaTool.getInputSchema().getProperties();
        List<String> requiredProperties = schemaTool.getInputSchema().getRequired();
        for (Parameter parameter : method.getParameters()) {
            Class<?> parameterType = parameter.getType();
            if (!isBasicType(parameterType)) {
                throw new McpServerException("Parameter Only Support Basic Type, Current Type[" + parameterType + "] Not Support.");
            }
            McpToolProperty mcpToolProperty = parameter.getAnnotation(McpToolProperty.class);
            String name = parameter.getName();
            ToolSchema.Property property = new ToolSchema.Property();
            if (mcpToolProperty != null) {
                if (!StringUtils.isBlank(mcpToolProperty.value())) {
                    name = mcpToolProperty.value();
                }
                if (!StringUtils.isBlank(mcpToolProperty.type())) {
                    property.setType(mcpToolProperty.type());
                }
                if (!StringUtils.isBlank(mcpToolProperty.description())) {
                    property.setDescription(mcpToolProperty.description());
                }
                if (mcpToolProperty.required()) {
                    requiredProperties.add(name);
                }
            } else {
                requiredProperties.add(name);
            }
            if (StringUtils.isEmpty(property.getType())) {
                property.setType(FitValueHelper.getInstance().getCommonBasicTypeName(parameterType));
            }
            properties.put(name, property);
        }
        McpToolMetadata toolMetadata = new McpToolMetadata();
        toolMetadata.setBean(bean);
        toolMetadata.setMethod(method);
        toolMetadata.setTool(tool);
        getMcpServerCapability().addToolMetadata(toolMetadata);
    }

    /**
     * 创建工具
     * @param method 方法对象
     * @return 工具对象
     */
    private Tool createTool(Method method) {
        McpTool mcpTool = method.getAnnotation(McpTool.class);
        ToolSchema.Tool schematool = getSchemaTool(mcpTool);
        Tool tool = new Tool();
        tool.setTool(schematool);
        tool.setExceptionHandler(getExceptionHandler(mcpTool.exceptionHandler()));
        return tool;
    }

    /**
     * 获取工具对象
     * @param mcpTool Mcp工具注解
     * @return 工具对象
     */
    private ToolSchema.Tool getSchemaTool(McpTool mcpTool) {
        ToolSchema.Tool schemaTool = new ToolSchema.Tool();
        schemaTool.setName(mcpTool.name());
        schemaTool.setDescription(mcpTool.description());
        ToolSchema.JsonSchema jsonSchema = new ToolSchema.JsonSchema();
        jsonSchema.setType("object");
        List<String> requiredProperties = new ArrayList<>();
        jsonSchema.setRequired(requiredProperties);
        Map<String, ToolSchema.Property> properties = new HashMap<>();
        jsonSchema.setProperties(properties);
        schemaTool.setInputSchema(jsonSchema);
        return schemaTool;
    }
}
